use sq_global::{API_STOCK_COMP, API_STOCK_DAILY};

use polars::prelude::{
    DataType, IntoLazy, RoundMode, SortMultipleOptions, col, lit,
};
use tokio::try_join;
use sq_macro::test_with_ctx;
use sq_polars::{get_col_idx, get_str_col};

#[test_with_ctx]
pub async fn test() {
    run().await.expect("run failed");
}

pub async fn run() -> anyhow::Result<Vec<String>> {
    let (com_df, daily_dfs) = try_join!(
        sq_tushare::load_data(API_STOCK_COMP, Vec::new()),
        sq_tushare::recent_dates(API_STOCK_DAILY, 10)
    )?;

    let daily_df = sq_polars::concat_dfs(daily_dfs)?;

    let latest_close = daily_df
        .select([col("ts_code"), col("trade_date"), col("pct_chg")])
        .sort(
            ["trade_date"],
            SortMultipleOptions::new().with_order_descending(true),
        )
        .group_by_stable(["ts_code"])
        .agg([col("trade_date").first(), col("pct_chg").first()]);

    let calc = com_df
        .lazy()
        .filter(col("province").is_not_null())
        .select([col("ts_code"), col("province")])
        .inner_join(latest_close, "ts_code", "ts_code")
        .group_by(["province"])
        .agg([
            col("trade_date").first(),
            col("pct_chg").mean().alias("pct_chg_avg"),
        ])
        .sort(
            ["pct_chg_avg"],
            SortMultipleOptions::new().with_order_descending(true),
        )
        .with_column(
            (col("pct_chg_avg")
                .round(2, RoundMode::HalfAwayFromZero)
                .cast(DataType::String)
                + lit("%"))
            .alias("pct_chg_avg_fmt"),
        )
        .collect()?;

    tracing::info!("{:?}", calc);

    let map = get_col_idx!(&calc);

    let province_idx = map["province"];
    let pct_chg_acg_fmt_idx = map["pct_chg_avg_fmt"];
    let trade_date_idx = map["trade_date"];

    let mut msg_all = Vec::new();
    for idx in 0..calc.height() {
        let row = calc.get_row(idx)?;
        let province = get_str_col!(row, province_idx);
        let pct_chg_acg_fmt = get_str_col!(row, pct_chg_acg_fmt_idx);
        let trade_date = get_str_col!(row, trade_date_idx);

        let msg = format!(
            "地区：{province} 的上市企业最近1个交易日{trade_date} 的平均涨幅为{pct_chg_acg_fmt}。\n\n"
        );
        msg_all.push(msg);
    }
    
    Ok(msg_all)
}
